package kz.kstu.aec.dsa;

import kz.kstu.aec.curve.EllipticCurveHelper;
import kz.kstu.aec.model.ECSignature;
import kz.kstu.aec.util.CryptographyUtil;
import sun.security.ec.ECPrivateKeyImpl;
import sun.security.ec.ECPublicKeyImpl;

import java.math.BigInteger;
import java.security.InvalidKeyException;
import java.security.spec.ECParameterSpec;
import java.security.spec.ECPoint;
import java.security.spec.EllipticCurve;
import java.util.Random;

/**
 * @author Filipp_Stankevich
 *         created: 27.01.12 19:18
 */
public class DSAHelper {

    public static ECPrivateKeyImpl createPrivateKey(BigInteger privateKey) {
        if (privateKey.signum() < 0) {
            throw new ArithmeticException("private key must be greater then zero");
        }
        ECParameterSpec curveSpec = kz.kstu.aec.curve.factory.EllipticCurveFactory.getEllipticCurveSpec();
        try {
            return new ECPrivateKeyImpl(privateKey, curveSpec);
        } catch (InvalidKeyException e) {
            throw new IllegalArgumentException("Internal error. Can not generate private key.");
        }
    }

    public static ECPrivateKeyImpl getPrivateKey(String login, String password) {
        byte[] hash = CryptographyUtil.getHash(login + password);
        BigInteger hashValue = new BigInteger(hash);
        BigInteger privateKey = hashValue.abs();
        return createPrivateKey(privateKey);
    }

    public static ECPublicKeyImpl generatePublicKey(ECPrivateKeyImpl privateKey) {
        ECParameterSpec curveSpec = privateKey.getParams();
        EllipticCurve curve = curveSpec.getCurve();
        ECPoint generator = curveSpec.getGenerator();
        ECPoint privatePoint = EllipticCurveHelper.multiplyPoints(generator, privateKey.getS(), curve);
        try {
            return new ECPublicKeyImpl(privatePoint, curveSpec);
        } catch (InvalidKeyException e) {
            throw new IllegalArgumentException("Internal error. Can not generate public key.");
        }
    }

    public static ECSignature sign(byte[] message, ECPrivateKeyImpl privateKey) {
        BigInteger r = BigInteger.ZERO;
        BigInteger s = BigInteger.ZERO;
        byte[] hash = CryptographyUtil.getHash(message);
        BigInteger h = new BigInteger(hash);
        Random random = new Random();
        while (BigInteger.ZERO.equals(r) || BigInteger.ZERO.equals(s)) {
            int rand = random.nextInt();
            while (rand <= 0) {
                rand = random.nextInt();
            }
            BigInteger k = BigInteger.valueOf(rand);
            ECParameterSpec curveSpec = privateKey.getParams();
            EllipticCurve curve = curveSpec.getCurve();
            ECPoint pPoint = EllipticCurveHelper.multiplyPoints(curveSpec.getGenerator(), k, curve);
            BigInteger x = pPoint.getAffineX();
            BigInteger cofactor = BigInteger.valueOf(curveSpec.getCofactor());
            BigInteger q = curveSpec.getOrder().divide(cofactor);
            r = x.mod(q);
            if (!BigInteger.ZERO.equals(r)) {
                s = k.multiply(h).add(r.multiply(privateKey.getS())).mod(q);
            }
        }
        return new ECSignature(r, s);
    }

    public static boolean checkSignature(ECSignature signature, ECPublicKeyImpl publicKey, byte[] message) {
        byte[] hash = CryptographyUtil.getHash(message);
        BigInteger h = new BigInteger(hash);
        ECParameterSpec curveSpec = publicKey.getParams();
        BigInteger cofactor = new BigInteger(String.valueOf(curveSpec.getCofactor()));
        BigInteger q = curveSpec.getOrder().divide(cofactor);
        BigInteger s = signature.getS();
        BigInteger r = signature.getR();
        if (BigInteger.ZERO.compareTo(r) >= 0 || s.compareTo(q) >= 0) {
            return false;
        }
        BigInteger u1 = s.multiply(h.modInverse(q)).mod(q);
        BigInteger u2 = r.multiply(BigInteger.valueOf(-1)).multiply(h.modInverse(q)).mod(q);
        EllipticCurve curve = curveSpec.getCurve();
        ECPoint generator = curveSpec.getGenerator();
        ECPoint firstPoint = EllipticCurveHelper.multiplyPoints(generator, u1, curve);
        ECPoint secondPoint = EllipticCurveHelper.multiplyPoints(publicKey.getW(), u2, curve);
        ECPoint pPoint = EllipticCurveHelper.addPoints(firstPoint, secondPoint, curve);
        if (ECPoint.POINT_INFINITY.equals(pPoint)) {
            return false;
        }
        BigInteger x = pPoint.getAffineX();
        BigInteger result = x.mod(q);
        return result.equals(r);
    }
}
